import torch
import matplotlib.pyplot as plt
#### 准备数据
x = torch.rand([50,1])
y = 5 * x + 30

### 初始化w,b
w = torch.rand(1,requires_grad=True)
b = torch.rand(1,requires_grad=True)

### 梯度下降
lr = 0.0001
for i in range(50000):
    y_pred = torch.matmul(x,w) + b
    loss = (y_pred - y).pow(2).mean()
    if w.grad is not None:
        w.grad.data.zero_()
    if b.grad is not None:
        b.grad.data.zero_()
    loss.backward()

    if i % 1000 == 0:
        print('epoch:%d,loss:%.7f'%(i,loss))

    w.data = w.data - lr * w.grad.data
    b.data = b.data - lr * b.grad.data

print('w:',w.data)
print('b:',b.data)
y_pred = torch.matmul(x,w) + b
plt.scatter(x.data,y.data,c = 'b')
plt.plot(x.data,y_pred.data,c = 'r')
plt.show()
